---
title: "Muris (2019), Section D.1, Design 1: DPD ADL(2), two cohorts four time periods"
output:
  html_document:
    df_print: paged
---

## Setup

```{r}
library(AER)
library(tidyverse)
library(cowplot)
library(gmm)
library(plm)
```

## Data generation

The function in the following chunk generates data for one person.

```{r}
DGP_dynamic_i <- function(i, T = 4, burn_in = 1,
                          rho = 0.9, beta1 = 1, beta2 = 0.0, sigma_a = 1, sigma_u = 0.1,
                          gamma0 = 1, gamma1 = 0.1, gamma2 = 0.4, tau = 1,sigma_v = 0.1){
  
  #############################################
  ### Explanation of the parameter values #####
  #############################################
  #
  ## general params
  # i is the personal identifier.
  # T is the number of time periods to export.
  # burn_in is a number of periods to get the process started (will be discarded)
  #
  ## outcome equation
  # rho <- 0.5 # coefficient on lagged Y
  # beta1 coefficient on xit
  # beta2 coefficient on xi,t-1
  # sigma_a is the variance of the unobserved heterogeneity
  # sigma_u is variance of error term in the outcomes equation
  #
  ## equation for x's
  # gamma0 is the constant term in this equation
  # gamma1 is the coefficient on first lag
  # gamma2 is the coefficient on the second lag
  # tau is the coefficient on alpha_i 
  # sigma_v variance of error terms in the equation for x's
  
  # Draw the parameter governing unobserved heterogeneity.
  alpha_i <- sigma_a*rnorm(1)
  # Startging value for the dependent variable
  Xi0 <- Ximin1 <- Yi0 <-  0
  # Set up the container.
  Y_i <- X_i <- numeric(burn_in+T)
  
  # Compute the first two periods: X
  X_i[1] <- gamma0  + gamma1*Xi0    + gamma2*Ximin1 + tau*alpha_i   + sigma_v*rnorm(1)
  X_i[2] <- gamma0  + gamma1*X_i[1] + gamma2*Xi0    + tau*alpha_i   + sigma_v*rnorm(1)
  # First two periods, T
  Y_i[1] <- alpha_i + rho*Yi0       + beta1*X_i[1] + beta2*Xi0     + sigma_u*rnorm(1)
  Y_i[2] <- alpha_i + rho*Y_i[1]    + beta1*X_i[2] + beta2*X_i[1]  + sigma_u*rnorm(1)
  
  # Remaining periods
  for(t in 3:(burn_in+T)){
    X_i[t] <- gamma0  + gamma1*X_i[t-1] + gamma2*X_i[t-2]  + tau*alpha_i     + sigma_v*rnorm(1)
    Y_i[t] <- alpha_i + rho*Y_i[t-1]    + beta1*X_i[t]     + beta2*X_i[t-1]  + sigma_u*rnorm(1)
  }
  # Return the final T periods
  export_period <- (burn_in+1):(burn_in+T)
  Y_i <- Y_i[export_period]
  X_i <- X_i[export_period]
  t <- 1:T
  return(as_tibble(cbind(Y_i,X_i,t,i)))
}
```

Check that this function works:

```{r}
DGP_dynamic_i(4)
```

## A DGP with two cohorts

Now we can generate a data set of size `n` from this function, using `map_df`. This corresponds to the first cohort, with cohort indicator `C=0`.

```{r}
n <- 100
cohort0 <- 1:n %>% map_df(DGP_dynamic_i) %>% mutate(C = 0)
```

Now generate data for another cohort. Note: For cohort 0, we used the default values `(gamma0 = 1, gamma1 = 0.1, gamma2 = 0.4, tau = 1,sigma_v = 0.1)`. Cohort 1, `C=1`, has different parameter values.

```{r}
cohort1 <- 1:n %>% map_df(DGP_dynamic_i, gamma0 = 1, gamma1 = 0, gamma2 = 0.8, tau = 1) %>% 
  mutate(
    C = 1, # guarantees that ids are unique once we merge cohorts
    i = i + n
  ) 
```

Put them together and turn into cross-section format.

```{r}
df_s <- rbind(cohort0,cohort1) %>% group_by(i) %>% 
  summarize(
    dY = Y_i[4] - Y_i[3],
    dX = X_i[4] - X_i[3],
    dLY = Y_i[3] - Y_i[2], 
    dLX = X_i[3] - X_i[2], 
    Y1 = Y_i[1],
    X1 = X_i[1],
    C = C[1])
```

Before programming the GMM estimator that follows exactly the format in the paper, try the IV regression corresponding to the efficient instruments:

```{r}
result1 <- ivreg(dY ~ -1 + dLY + dX + dLX | Y1 + X1 + C:X1 + C:Y1 , data = df_s)
summary(result1)
```

This suggests that the DGP functions are OK.

## Efficient and oracle estimators via the `gmm` package

The `gmm` package will do the heavy lifting. All we need to do is program the moment function, one for each estimator.

This first function corresponds to the infeasible estimator that corresponds to the case of complete data. It is useful as a benchmark. 

```{r}
psi <- function(theta,Z){
  
  n <- nrow(Z)
  p <- 6
  
  # Load the parameters
  rho = theta[1]
  beta1 = theta[2]
  beta2 = theta[3]
  
  # Compute the differenced residuals
  du5 = Z[,"dY5"] - rho*Z[,"dY4"] - beta1*Z[,"dX5"] - beta2*Z[,"dX4"]
  du4 = Z[,"dY4"] - rho*Z[,"dY3"] - beta1*Z[,"dX4"] - beta2*Z[,"dX3"]
  
  return( 
    cbind(
      Z[,"Y2"] * du5,
      Z[,"Y1"] * du5,
      Z[,"X2"] * du5,
      Z[,"X1"] * du5,
      # 
      Z[,"Y1"] * du4,
      Z[,"X1"] * du4
    )
  )  
}
```

Set up the two-cohort moments that are defined for the incomplete data, `psiJ`. This corresponds to the efficient estimator in the paper.

```{r}
psiJ <- function(theta,Z){
  
  n <- nrow(Z)
  
  # Load the parameters
  rho = theta[1]
  beta1 = theta[2]
  beta2 = theta[3]
  
  # Compute the differenced residuals
  du5 = Z[,"dY5"] - rho*Z[,"dY4"] - beta1*Z[,"dX5"] - beta2*Z[,"dX4"]
  du4 = Z[,"dY4"] - rho*Z[,"dY3"] - beta1*Z[,"dX4"] - beta2*Z[,"dX3"]
  
  return( 
    cbind(
      # For cohort == 1, the group that starts in period 2 and goes through period 5
      ifelse(Z[,"cohort"] == 1,1,0) * Z[,"Y2"] * du5,
      ifelse(Z[,"cohort"] == 1,1,0) * Z[,"X2"] * du5,
      # For cohort == 0, the group that starts in period 1 and goes through period 4.
      ifelse(Z[,"cohort"] == 0,1,0) * Z[,"Y1"] * du4,
      ifelse(Z[,"cohort"] == 0,1,0) * Z[,"X1"] * du4
    )
  )  
}
```

Set up the DGP in line with Section D.1 in Muris (2019), and apply both estimators.

```{r}
n <- 100
cohort0 <- 1:n %>% map_df(DGP_dynamic_i, T = 5, burn_in = 1,
                          rho = 0.9, beta1 = 0.5, beta2 = 0.2,
                          gamma0 = 1, gamma1 = 0.3, gamma2 = 0.3, tau = 0.5,sigma_v = 0.1) %>% 
  mutate(C = 0)

cohort1 <- 1:n %>% map_df(DGP_dynamic_i, T = 5, burn_in = 1,
                          rho = 0.9, beta1 = 0.5, beta2 = 0.2,
                          gamma0 = 2, gamma1 = 0.6, gamma2 = 0.1, tau = 0.8,sigma_v = 0.2) %>% 
  mutate(
    C = 1,
    i = i + n) 

df_new <- rbind(cohort0,cohort1) %>% group_by(i) %>% 
  summarize(
    
    cohort = C[1],
    
    Y1 = Y_i[1],
    Y2 = Y_i[2],
    X1 = X_i[1],
    X2 = X_i[2],
    
    dY5 = Y_i[5]-Y_i[4],
    dY4 = Y_i[4]-Y_i[3],
    dY3 = Y_i[3]-Y_i[2],
    
    dX5 = X_i[5]-X_i[4],
    dX4 = X_i[4]-X_i[3],
    dX3 = X_i[3]-X_i[2],
    dX2 = X_i[2]-X_i[1]
  )

out1 <- gmm(g = psi, x = as.matrix(df_new), t0 = c(0.9,0.5,0.2))
out2 <- gmm(g = psiJ, x = as.matrix(df_new), t0 = c(0.9,0.5,0.2))
```

Now look at the results.

```{r}
cbind(out1$coefficients,out2$coefficients)
```

It seems to be working! Wrap the procedure above in a function that represents one simulation.

```{r}
one_sim <- function(s, n = 100, rho_0 = 0.9, beta1_0 = 0.5, beta2_0 = 0.2, D_D_X = 0.1) {
  
  cohort0 <- 1:n %>% map_df(DGP_dynamic_i, T = 5, burn_in = 1,
                            rho = rho_0, beta1 = beta1_0, beta2 = beta2_0,
                            gamma0 = 1,         gamma1 = 0.4,         gamma2 = 0.4,         tau = 0.4,        sigma_v = 0.1) %>% 
    mutate(C = 0)
  
  cohort1 <- 1:n %>% map_df(DGP_dynamic_i, T = 5, burn_in = 1,
                            rho = rho_0, beta1 = beta1_0, beta2 = beta2_0,
                            gamma0 = 1 + D_D_X, gamma1 = 0.4 + D_D_X, gamma2 = 0.4 - D_D_X, tau = 0.4 + D_D_X,sigma_v = 0.1) %>% 
    mutate(
      C = 1,
      i = i + n) # guarantees that ids are unique
  
  df_new <- rbind(cohort0,cohort1) %>% group_by(i) %>% 
    summarize(
      
      cohort = C[1],
      
      Y1 = Y_i[1],
      Y2 = Y_i[2],
      X1 = X_i[1],
      X2 = X_i[2],
      
      dY5 = Y_i[5]-Y_i[4],
      dY4 = Y_i[4]-Y_i[3],
      dY3 = Y_i[3]-Y_i[2],
      
      dX5 = X_i[5]-X_i[4],
      dX4 = X_i[4]-X_i[3],
      dX3 = X_i[3]-X_i[2],
      dX2 = X_i[2]-X_i[1]
    )
  
  out1 <- gmm(g = psi, x = as.matrix(df_new),  t0 = c(rho_0,beta1_0,beta2_0))
  out2 <- gmm(g = psiJ, x = as.matrix(df_new), t0 = c(rho_0,beta1_0,beta2_0))
  out1_sum <- summary(out1)
  out2_sum <- summary(out2)
  
  return(
    tibble(
      s = c(s,s),
      
      rho_0 = rho_0,
      D_D_X = D_D_X,
      
      rho = c(out1$coefficients[1],out2$coefficients[1]),
      rho_dev = rho - rho_0,
      rho_se = c(out1_sum$coefficients[,"Std. Error"][1],
                 out2_sum$coefficients[,"Std. Error"][1]),
      rho_lb = c(confint.gmm(out1)$test[1,1],confint.gmm(out2)$test[1,1]),
      rho_ub = c(confint.gmm(out1)$test[1,2],confint.gmm(out2)$test[1,2]),
      
      beta1 = c(out1$coefficients[2],out2$coefficients[2]),
      beta1_se = c(out1_sum$coefficients[,"Std. Error"][2],
                 out2_sum$coefficients[,"Std. Error"][2]),
      beta1_lb = c(confint.gmm(out1)$test[2,1],confint.gmm(out2)$test[2,1]),
      beta1_ub = c(confint.gmm(out1)$test[2,2],confint.gmm(out2)$test[2,2]),

      beta2 = c(out1$coefficients[3],out2$coefficients[3]),
      beta2_se = c(out1_sum$coefficients[,"Std. Error"][3],
                 out2_sum$coefficients[,"Std. Error"][3]),
      beta2_lb = c(confint.gmm(out1)$test[3,1],confint.gmm(out2)$test[3,1]),
      beta2_ub = c(confint.gmm(out1)$test[3,2],confint.gmm(out2)$test[3,2]),
      
      method = c("full","incomplete")
    )
  )
  
}
```

Test-run using 10 simulations.

```{r}
result <- 1:10 %>% map_df(one_sim)
result
```

The following function then corresponds to a simulation study with `S` simulations.

```{r}
sim_study <- function(S,n = 100, rho_0 = 0.9, D_D_X = 0.3){
  
  # Need the betas (baked into `one_sim`)
  beta1_0 = 0.5
  beta2_0 = 0.2
  
  # Run it!
  result <- 1:S %>% map_df(one_sim, n = n, rho_0 = rho_0, D_D_X = D_D_X)

  # Process it
  result_table <- result %>% 
    mutate(
      rho_in =   ifelse(rho_0 >= rho_lb,1,0) * ifelse(rho_0 <= rho_ub,1,0),
      beta1_in = ifelse(beta1_0 >= beta1_lb,1,0) * ifelse(beta1_0 <= beta1_ub,1,0),
      beta2_in = ifelse(beta2_0 >= beta2_lb,1,0) * ifelse(beta2_0 <= beta2_ub,1,0)
    ) %>% 
    group_by(method) %>% 
    summarize(
      
      rho_0 = rho_0[1],
      D_D_X = D_D_X[1],
      
      # The autoregressive parameter.
      MSE_rho = sqrt(mean((rho - rho_0)^2)),
      bias_rho = mean(rho - rho_0),
      SE_rho = sd(rho),
      meanSE_rho = mean(rho_se),
      rho_length = mean(rho_ub - rho_lb),
      rho_coverage = mean(rho_in),
      
      # Parameter on x_{it}
      MSE_beta1 = sqrt(mean((beta1 - beta1_0)^2)),
      bias_beta1 = mean(beta1 - beta1_0),
      SE_beta1 = sd(beta1),
      meanSE_beta1 = mean(beta1_se),
      beta1_length = mean(beta1_ub - beta1_lb),
      beta1_coverage = mean(beta1_in),
      
      # Parameter on x_{i,t-1}
      MSE_beta2 = sqrt(mean((beta2 - beta2_0)^2)),
      bias_beta2 = mean(beta2 - beta2_0),
      SE_beta2 = sd(beta2),
      meanSE_beta2 = mean(beta2_se),
      beta2_length = mean(beta2_ub - beta2_lb),
      beta2_coverage = mean(beta2_in)
      
    )
  
  # Return it
  return(list(result = result, result_table = result_table))
    
}
```

## Result generation

The following chunk contains 6 simulation studies, one for each parameter combination considered in Table 5.

```{r}
SS <- 1000
nn <- 100

ss_0701 <- sim_study(S = SS, n = nn, rho_0 = 0.7, D_D_X = 0.1)
ss_0703 <- sim_study(S = SS, n = nn, rho_0 = 0.7, D_D_X = 0.3)
ss_0801 <- sim_study(S = SS, n = nn, rho_0 = 0.8, D_D_X = 0.1)
ss_0803 <- sim_study(S = SS, n = nn, rho_0 = 0.8, D_D_X = 0.3)
ss_0901 <- sim_study(S = SS, n = nn, rho_0 = 0.9, D_D_X = 0.1)
ss_0903 <- sim_study(S = SS, n = nn, rho_0 = 0.9, D_D_X = 0.3)
```

Now gather the results for Figure 1 (fix `D_D_X=0.3`) in a `tibble` and generate the plot corresponding to Figure D.1.

```{r}
result_design1 <- rbind(
  ss_0703$result,
  ss_0803$result,
  ss_0903$result
)

(fig_design_1 <- result_design1 %>% ggplot() + geom_boxplot(aes(x = method, y = rho)) + facet_grid(~rho_0) + theme(axis.text.x = element_text(angle=45,vjust=0.5)) + xlab(""))
```

Save Figure D.1 as a pdf.

```{r}
save_plot("dpd_design_1.pdf",fig_design_1)
```

Finally, gather all results to generate Table 5.

```{r}
table_design1 <- rbind(
  ss_0701$result_table,
  ss_0801$result_table,
  ss_0901$result_table,
  
  ss_0703$result_table,
  ss_0803$result_table,
  ss_0903$result_table
)

table_design1
```

Save this collection of results in case you want to create additional summaries, and don't want to run the code again.

```{r}
save.image("dpd_cohorts_design1.RData")
```
